/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia.translation.functions;



import com.bff.gaia.unified.runners.core.construction.SerializablePipelineOptions;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.transforms.CombineFnBase;

import com.bff.gaia.unified.sdk.transforms.Combine;

import com.bff.gaia.unified.sdk.transforms.windowing.BoundedWindow;

import com.bff.gaia.unified.sdk.transforms.windowing.IntervalWindow;

import com.bff.gaia.unified.sdk.util.WindowedValue;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PCollectionView;

import com.bff.gaia.unified.sdk.values.WindowingStrategy;

import com.bff.gaia.api.common.functions.RichGroupReduceFunction;

import com.bff.gaia.util.Collector;



import java.util.Map;



/**

 * This is the second part for executing a {@link Combine.PerKey} on

 * Gaia, the second part is {@link GaiaReduceFunction}. This function performs the final

 * combination of the pre-combined values after a shuffle.

 *

 * <p>The input to {@link #reduce(Iterable, Collector)} are elements of the same key but for

 * different windows. We have to ensure that we only combine elements of matching windows.

 */

public class GaiaReduceFunction<K, AccumT, OutputT, W extends BoundedWindow>

    extends RichGroupReduceFunction<WindowedValue<KV<K, AccumT>>, WindowedValue<KV<K, OutputT>>> {



  protected final CombineFnBase.GlobalCombineFn<?, AccumT, OutputT> combineFn;



  protected final WindowingStrategy<Object, W> windowingStrategy;



  // TODO: Remove side input functionality since liftable Combines no longer have side inputs.

  protected final Map<PCollectionView<?>, WindowingStrategy<?, ?>> sideInputs;



  protected final SerializablePipelineOptions serializedOptions;



  public GaiaReduceFunction(

      CombineFnBase.GlobalCombineFn<?, AccumT, OutputT> combineFn,

      WindowingStrategy<Object, W> windowingStrategy,

      Map<PCollectionView<?>, WindowingStrategy<?, ?>> sideInputs,

      PipelineOptions pipelineOptions) {



    this.combineFn = combineFn;



    this.windowingStrategy = windowingStrategy;

    this.sideInputs = sideInputs;



    this.serializedOptions = new SerializablePipelineOptions(pipelineOptions);

  }



  @Override

  public void reduce(

      Iterable<WindowedValue<KV<K, AccumT>>> elements, Collector<WindowedValue<KV<K, OutputT>>> out)

      throws Exception {



    PipelineOptions options = serializedOptions.get();



    GaiaSideInputReader sideInputReader =

        new GaiaSideInputReader(sideInputs, getRuntimeContext());



    AbstractGaiaCombineRunner<K, AccumT, AccumT, OutputT, W> reduceRunner;



    if (!windowingStrategy.getWindowFn().isNonMerging()

        && !windowingStrategy.getWindowFn().windowCoder().equals(IntervalWindow.getCoder())) {

      reduceRunner = new HashingGaiaCombineRunner<>();

    } else {

      reduceRunner = new SortingGaiaCombineRunner<>();

    }

    reduceRunner.combine(

        new AbstractGaiaCombineRunner.FinalGaiaCombiner<>(combineFn),

        windowingStrategy,

        sideInputReader,

        options,

        elements,

        out);

  }

}